import unittest
from src.dataset import Vocabulary


class TestVocabulary(unittest.TestCase):
    def setUp(self):
        self.vocab = Vocabulary()

    def test_add_sentence(self):
        self.vocab.add_sentence("hello world")

        # 确保单词正确添加（小写）
        self.assertIn("hello", self.vocab.word2idx)
        self.assertIn("world", self.vocab.word2idx)
        self.assertEqual(self.vocab.word_count["hello"], 1)
        self.assertEqual(self.vocab.word_count["world"], 1)

    def test_build(self):
        sentences = ["hello world", "hello hello world"]
        self.vocab.build(sentences, min_freq=2)

        # "hello" 频次 3，"world" 频次 2，应该都被保留
        self.assertIn("hello", self.vocab.word2idx)
        self.assertIn("world", self.vocab.word2idx)
        self.assertEqual(len(self.vocab.word2idx), 6)  # 4 个特殊标记 + 2 个符合条件的单词

    def test_index_mapping(self):
        self.vocab.add_sentence("new word")
        idx_new = self.vocab.word2idx["new"]
        idx_word = self.vocab.word2idx["word"]

        # 确保索引和单词匹配
        self.assertEqual(self.vocab.idx2word[idx_new], "new")
        self.assertEqual(self.vocab.idx2word[idx_word], "word")